import numpy as np
import matplotlib.pyplot as plt
import pylab
import random
from numpy import sin, arccos, cos, pi, array, zeros, sqrt, exp

### calculate the mPDF function f(r)
### r is the position array, J is the spin quantum number, spinlocations is an array containing the positions of all the spins, spinvectors is an array giving the spins in the same sequence as the spinlocations array
def frcalculator(r,J,spinlocations,spinvectors):
    '''Calculates the mPDF given lists of magnetic moments and their positions.

Returns the mPDF function f(r) as a numpy array of the same shape as the input
array r.

Parameters
----------
r : array_like
    Array containing the r-points for which the mPDF f(r) should be computed.
J : float
    The total angular momentum quantum number of the system.
spinlocations : array_like
    An array of (x,y,z) tuples giving the real-space positions of all the 
    magnetic atoms in the system.
spinvectors: array_like
    An array of (x,y,z) tuples giving the magnetic moments of all the magnetic
    atoms in the system, in the same order as spinlocations.

Returns
-------
f : array_like r
    The real-space mPDF function f(r).
'''
    f=np.zeros_like(r)
    for i in range(len(spinvectors)):
        print 'Working on: '+str(i)+'/'+str(len(spinvectors))
        si,ri=spinvectors[i],spinlocations[i]
        otherspinvecs=np.delete(spinvectors,i,0)
        otherspinlocs=np.delete(spinlocations,i,0)
        for j in range(len(otherspinvecs)):
            sj,rj=otherspinvecs[j],otherspinlocs[j]
            rij=np.linalg.norm(rj-ri)
            if rij <= r[-1]:
                C1=(si[0]*(rj[0]-ri[0])+si[1]*(rj[1]-ri[1])+si[2]*(rj[2]-ri[2]))/rij**2
                try:
                    C2sq=1/((si[0]-C1*(rj[0]-ri[0]))**2+(si[1]-C1*(rj[1]-ri[1]))**2+(si[2]-C1*(rj[2]-ri[2]))**2)
                    aij=C2sq*(si[0]*(si[0]-C1*(rj[0]-ri[0]))+si[1]*(si[1]-C1*(rj[1]-ri[1]))+si[2]*(si[2]-C1*(rj[2]-ri[2])))*(sj[0]*(si[0]-C1*(rj[0]-ri[0]))+sj[1]*(si[1]-C1*(rj[1]-ri[1]))+sj[2]*(si[2]-C1*(rj[2]-ri[2])))
                except: ### if spin si is aligned with vector rj-ri, then the local y-axis is ill-defined and aij is zero.
                    aij=0
                bij=2*(1/rij**2)*(si[0]*(rj[0]-ri[0])+si[1]*(rj[1]-ri[1])+si[2]*(rj[2]-ri[2]))*(sj[0]*(rj[0]-ri[0])+sj[1]*(rj[1]-ri[1])+sj[2]*(rj[2]-ri[2]))-aij
                cutoff=(np.abs(r-rij)).argmin()
                f[:cutoff]+=bij*r[:cutoff]/rij**3
                if rij<(r[-1]+.001):
                    f[cutoff]+=aij/rij+0.5*bij/rij**2
    f=(3./(len(spinlocations)*2*J*(J+1)))*f
    return f

### Faster calculator that does not include the sum over each individual spin at the origin. Here I just choose the first spin to be at the origin. 
def frcalculatorshort(r,J,spinlocations,spinvectors):
	f=np.zeros_like(r)
	i=0
	si,ri=spinvectors[i],spinlocations[i]
	otherspinvecs=np.delete(spinvectors,i,0)
	otherspinlocs=np.delete(spinlocations,i,0)
	for j in range(len(otherspinvecs)):
		sj,rj=otherspinvecs[j],otherspinlocs[j]
		rij=np.linalg.norm(rj-ri)
		[six,siy]=localcoords(ri,rj,si,sj)[:2]
		[sjx,sjy]=localcoords(ri,rj,si,sj)[2:]
		aij=siy*sjy
		bij=2*six*sjx-aij
		cutoff=(np.abs(r-rij)).argmin()
		f[:cutoff]+=bij*r[:cutoff]/rij**3
		if rij<(r[-1]+.001):
			f[cutoff]+=aij/rij+0.5*bij/rij**2
		print str(j)+'/'+str(len(otherspinvecs))
	f=(3./(len(spinlocations)*2*J*(J+1)))*f
	return f

### This version uses the global coordinate system rather than a local coordinate system for each spin pair. Notation is uglier, but it works about 40% faster.	
def frcalculatorshort2(r,J,spinlocations,spinvectors):
	f=np.zeros_like(r)
	i=0
	si,ri=spinvectors[i],spinlocations[i]
	otherspinvecs=np.delete(spinvectors,i,0)
	otherspinlocs=np.delete(spinlocations,i,0)
	for j in range(len(otherspinvecs)):
		sj,rj=otherspinvecs[j],otherspinlocs[j]
		rij=np.linalg.norm(rj-ri)
		C1=(si[0]*(rj[0]-ri[0])+si[1]*(rj[1]-ri[1])+si[2]*(rj[2]-ri[2]))/rij**2
		try:
			C2sq=1/((si[0]-C1*(rj[0]-ri[0]))**2+(si[1]-C1*(rj[1]-ri[1]))**2+(si[2]-C1*(rj[2]-ri[2]))**2)
			aij=C2sq*(si[0]*(si[0]-C1*(rj[0]-ri[0]))+si[1]*(si[1]-C1*(rj[1]-ri[1]))+si[2]*(si[2]-C1*(rj[2]-ri[2])))*(sj[0]*(si[0]-C1*(rj[0]-ri[0]))+sj[1]*(si[1]-C1*(rj[1]-ri[1]))+sj[2]*(si[2]-C1*(rj[2]-ri[2])))
		except: ### if spin si is aligned with vector rj-ri, then the local y-axis is ill-defined and aij is zero.
			aij=0
		bij=2*(1/rij**2)*(si[0]*(rj[0]-ri[0])+si[1]*(rj[1]-ri[1])+si[2]*(rj[2]-ri[2]))*(sj[0]*(rj[0]-ri[0])+sj[1]*(rj[1]-ri[1])+sj[2]*(rj[2]-ri[2]))-aij
		cutoff=(np.abs(r-rij)).argmin()
		f[:cutoff]+=bij*r[:cutoff]/rij**3
		if rij<(r[-1]+.001):
			f[cutoff]+=aij/rij+0.5*bij/rij**2
		print str(j)+'/'+str(len(otherspinvecs))
	f=(3./(len(spinlocations)*2*J*(J+1)))*f
	return f	

### define the local coordinate system between two spins used for the calculation of the mPDF
def localcoords(r1,r2,s1,s2):
	xhat=(r2-r1)/np.linalg.norm(r2-r1)
	yhat=(s1-xhat*(np.dot(s1,xhat)))/np.linalg.norm(s1-xhat*(np.dot(s1,xhat)))
#	zhat=np.cross(xhat,yhat)
	return [np.dot(s1,xhat),np.dot(s1,yhat),np.dot(s2,xhat),np.dot(s2,yhat)]
	
### create a unit vector in a random direction
def unitvec():
	phi=random.uniform(0,2*pi)
	theta=arccos(random.uniform(-1,1))
	vector=np.array([cos(phi)*sin(theta),sin(phi)*sin(theta),cos(theta)])
	return vector
	
### calculate the orientational terms aij and bij for two spins
def abcalculator(ri,si,rj,sj):
	rij=np.linalg.norm(rj-ri)
	C1=(si[0]*(rj[0]-ri[0])+si[1]*(rj[1]-ri[1])+si[2]*(rj[2]-ri[2]))/rij**2
	try:
		C2sq=1/((si[0]-C1*(rj[0]-ri[0]))**2+(si[1]-C1*(rj[1]-ri[1]))**2+(si[2]-C1*(rj[2]-ri[2]))**2)
		aij=C2sq*(si[0]*(si[0]-C1*(rj[0]-ri[0]))+si[1]*(si[1]-C1*(rj[1]-ri[1]))+si[2]*(si[2]-C1*(rj[2]-ri[2])))*(sj[0]*(si[0]-C1*(rj[0]-ri[0]))+sj[1]*(si[1]-C1*(rj[1]-ri[1]))+sj[2]*(si[2]-C1*(rj[2]-ri[2])))
	except: ### if spin si is aligned with vector rj-ri, then the local y-axis is ill-defined and aij is zero.
		aij=0
	bij=2*(1/rij**2)*(si[0]*(rj[0]-ri[0])+si[1]*(rj[1]-ri[1])+si[2]*(rj[2]-ri[2]))*(sj[0]*(rj[0]-ri[0])+sj[1]*(rj[1]-ri[1])+sj[2]*(rj[2]-ri[2]))-aij	
	return [aij,bij]
